# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

# Helper functions for 3D sparse pattern
# These function are not optimized and very inefficient.
# Avoid calling them too frequent or use a cache mechanism.

from functools import lru_cache

import numpy as np
import torch
import torch_npu
# from vllm.triton_utils import triton
import triton

'''
获取输入数组中的有效数据
在data中记录相关所有的有效数据
indices中记录每一行中有效数据下标
indptr中记录每一行有效数据的长度
compressed sparse row

examples
tensor([[ True, False,  True, False],  # 第0行：位置0和2有效
        [False,  True,  True, False],  # 第1行：位置1和2有效
        [ True, False, False,  True]]) # 第2行：位置0和3有效

self.data = np.array([True, True, True, True])（4 个 True）
self.indices = np.array([0, 2, 1, 2, 0, 3])（所有有效列索引）
self.indptr = np.array([0, 2, 4, 6])（行指针：第0行[0:2]、第1行[2:4]、第2行[4:6]）
'''
class csr_matrix:
    """Simple implementation of CSR matrix conversion without scipy.
    This replaced scipy.sparse.csr_matrix() previously used."""

    def __init__(self, input_array):
        if not isinstance(input_array, np.ndarray):
            raise ValueError("Input must be a NumPy array")

        self.shape = input_array.shape
        rows, cols = self.shape
        # 存储有效值
        data = []
        # 存储列索引下标
        indices = []
        # 存储行指针 从0开始
        # 有效数据的长度
        indptr = [0]

        for i in range(rows):
            for j in range(cols):
                if input_array[i, j]:
                    data.append(input_array[i, j])
                    indices.append(j)
            indptr.append(len(indices))

        self.data = np.array(data)
        self.indices = np.array(indices)
        self.indptr = np.array(indptr)
        # CSR格式构建完成，indptr[i]:indptr[i+1]表示第i行的有效索引范围
'''
将2D或3D Pytorch张量转成CSR格式的行索引的（crow）和列索引（cols）
用于将稠密注意力掩码压缩为稀疏表示
这里是将dense数据转成compressed数据
'''
def dense_to_crow_col(x: torch.Tensor):
    """Turning a 2D/3D torch tensor (x) to CSR rows/cols indexing.
    NOTE: col_indices padded -1
    """
    device = x.device
    pad = -1
    dim = x.dim()
    # 只处理2，3维数据
    assert x.dim() in (2, 3)
    # 如果只有两维度 进行拓展一个维度
    if x.dim() == 2:
        x = x[None] # 相当于变成了(1, a, b) 在第0维进行拓展一个维度
    # 针对每个切片xi使用csr_matrix转成csr对象列表
    '''
    x的形状是(B,R,C) B是批次大小 for xi in x
    是逐个批次切片xi(R,C)转换成CSR对象
    '''
    x = [csr_matrix(xi.bool().cpu().numpy()) for xi in x]
    # vstack操作，将indptr数组转成行索引
    '''
    现在的x全是CSR对象的集合体，这里提取indptr对象，之后垂直堆叠成crows
    行指针索引内容
    '''
    crows = torch.vstack([torch.from_numpy(xi.indptr) for xi in x])
    # 将csr的indices建立列索引列表
    '''
    从每个CSR的indices创建列索引列表
    '''
    cols = [torch.from_numpy(xi.indices) for xi in x]
    # 计算最大列索引长度
    max_cols = max(len(xi) for xi in cols)
    # 对每个列索引，用-1填充到最大长度 这个部分进行堆叠处理
    '''
    遍历cols列索引列表 进行填充长度
    pad = -1 zeros 取零操作 通过pad + xi.zeros来全部填充-1值到max_cols
    '''
    cols = [
        torch.cat([xi, pad + xi.new_zeros(max_cols - xi.shape[0])])
        for xi in cols
    ]
    # 垂直堆叠填充后的列索引
    cols = torch.vstack(cols)
    # 后面将维度又降低回来了
    # 这里维度降回来了
    if dim == 2:
        crows = crows[0]
        cols = cols[0]
    return crows.to(device), cols.to(device)


'''
从csr格式的行索引（crows）和列索引（cols）重建稠密张量
这里是从compressed重建回dense
'''
def crow_col_to_dense(crows: torch.Tensor,
                      cols: torch.Tensor,
                      dtype: torch.dtype = torch.float16):
    dim = crows.dim()
    # 维度为1
    if dim == 1:
        # 进行拓展维度 拓展维度到[1, a]
        crows = crows[None]
        cols = cols[None]
    device = crows.device
    crows, cols = crows.cpu(), cols.cpu()  # faster in cpu
    # shape (batch 批次信息 rows 行数 cols 列数)
    shape = (crows.shape[0], crows.shape[1] - 1, cols.max() + 1)
    x = torch.zeros(shape, dtype=dtype)
    for i in range(shape[0]):
        for j in range(shape[1]):
            # 将一定范围内容的数据标记成有效数据
            x[i, j, cols[i, crows[i, j]:crows[i, j + 1]]] = 1
    if dim == 1:
        x = x[0]
    return x.to(device)

'''
类似于dense_to_crow_col，但是转换成了CSC（Compressed Sparse Column）格式
通过转置transpose操作，将行列进行转换操作
'''
def dense_to_ccol_row(x: torch.Tensor):
    """Similar, but to CSC format"""
    x = x.transpose(-2, -1)
    return dense_to_crow_col(x)


'''
从csc格式的列索引和行索引重建稠密张量
通过调用csr重建后转置维度，实现csc的反向转换
permute操作，将后面两个维度进行转换了
'''
def ccol_row_to_dense(ccol: torch.Tensor,
                      rows: torch.Tensor,
                      dtype: torch.dtype = torch.float16):
    return crow_col_to_dense(ccol, rows, dtype).permute(0, 2, 1).contiguous()


'''
用于生成同质头（homo_head）的稀疏注意力块掩码
基于块大小、局部块数和垂直步长，创建因果+稀疏的块掩码，支持返回稠密版本
'''
def _get_sparse_attn_mask_homo_head(
    q_len: int,
    max_seqlen: int,
    dtype: torch.dtype,
    device: torch.device,
    block_size: int = 128,
    local_blocks: int = 4,
    vert_stride: int = 4,
    return_dense: bool = False,
):
    """
    :return: a tuple of 3:
        - tuple of crow_indices, col_indices representation
            of CSR format.
        - block dense mask
        - all token dense mask (be aware that it can be
            OOM if it is too big) if `return_dense==True`,
            otherwise, None
    """
    with torch.no_grad():
        # 计算总块数
        num_blocks = triton.cdiv(max_seqlen, block_size)
        # 查询块位置 形状(num_blocks, 1)
        q_pos = torch.arange(num_blocks)[:, None]
        # 键块位置 形状(1, num_blocks)
        k_pos = torch.arange(num_blocks)[None]
        # 垂直步长掩码，每个vert_stride块选一
        mask_vert_strided = (torch.arange(num_blocks) + 1) % vert_stride == 0
        # 获取块掩码
        block_mask_dense = (((q_pos >= k_pos)
                             & ((q_pos - k_pos < local_blocks)
                                | mask_vert_strided)).to(device).to(dtype))
        # 查询块数
        num_blocks_q = triton.cdiv(q_len, block_size)
        # 获取最后num_blocks_q行数据，并转成csr格式
        block_mask_dense_output = (dense_to_crow_col(
            block_mask_dense[-num_blocks_q:].contiguous()))
    if return_dense:
        '''
        torch.kron操作的样例
        矩阵A:
        tensor([[1, 2],
                [3, 4]])
        矩阵B:
        tensor([[0, 5],
                [6, 7]])
        A和B的克罗内积C:
        tensor([[ 0,  5,  0, 10],
                [ 6,  7, 12, 14],
                [ 0, 15,  0, 20],
                [18, 21, 24, 28]])
        '''
        # 将kron的积拓展到token级
        mask_dense = torch.kron(
            block_mask_dense,
            block_mask_dense.new_ones((block_size, block_size)),
        )
        # 取下三角因果掩码，取最后q_len行 
        causal_mask = torch.tril(torch.ones(
            max_seqlen, max_seqlen)).type_as(mask_dense)[-q_len:]
        # 截取并乘上因果掩码
        mask_dense = mask_dense[-q_len:, :max_seqlen] * causal_mask
        return (
            block_mask_dense_output,
            block_mask_dense,
            mask_dense,
        )
    else:
        return (
            block_mask_dense_output,
            block_mask_dense,
            None,
        )


'''
将二进制掩码（0/1）转换为注意力偏置（bias）格式：将0（跳过）转为-inf（负无穷，屏蔽注意力），1保持为0
'''
def binary_mask_to_bias(mask_dense: torch.Tensor):
    # 这里将原来0的位置，填充为1
    mask_dense = 1 - mask_dense
    # 将原来0的位置填充为-inf，屏蔽注意力
    mask_dense.masked_fill_(mask_dense.bool(), -torch.inf)
    return mask_dense


'''
计算每个注意力头的垂直步长滑动偏移。
如果是同质头，返回0；否则，根据总头数和垂直步长均匀分配步长，确保多头稀疏覆盖不同垂直位置。
home_head是指维度、计算逻辑完全一致  important!
'''
def get_head_sliding_step(n_heads: int,
                          vert_stride: int,
                          homo_head: bool = False):
    if homo_head:
        return 0
    return max(1, int(vert_stride / n_heads))


'''
主函数，生成稀疏注意力掩码，支持同质头（所有头相同掩码）和非同质头（每个头滑动偏移）。
返回CSR格式块掩码、块级稠密掩码，以及可选的token级稠密掩码（binary或bias类型）。

homo_head = True 同质头，所有头共享相同掩码

'''
# 使用lru_cache缓存结果以提高效率。避免重复计算
@lru_cache
def get_sparse_attn_mask(
    n_heads: int,
    q_len: int,
    max_seqlen: int,
    dtype: torch.dtype,
    device: torch.device,
    block_size: int = 64,
    local_blocks: int = 4,
    vert_stride: int = 4,
    homo_head: bool = True,
    return_dense: bool = False,
    dense_mask_type: str = "binary",
):
    """
    :param dense_mask_type: "binary" (0 for skip token, 1 for others)
        or "bias" (-inf for skip token, 0 or others)
    :return: a tuple of 3:
        - tuple of crow_indices, col_indices representation
            of CSR format.
        - block dense mask
        - all token dense mask (be aware that it can be OOM if it
            is too big) if `return_dense==True`, otherwise, None
    """
    assert dense_mask_type in ("binary", "bias")
    # 同质头操作
    # homo_head = True 是单掩码复制
    # homo_head = False 是逐头生成
    if homo_head:
        with torch.no_grad():
            (crow, col), block_mask_dense, mask_dense = (
                _get_sparse_attn_mask_homo_head(
                    q_len,
                    max_seqlen,
                    dtype,
                    device,
                    block_size,
                    local_blocks,
                    vert_stride,
                    return_dense,
                ))
            # 先进行拓展一个维度，之后拓展到n_heads长度
            # 通过expand操作，将crow和col复制到所有头上
            # 形状从(1, rows + 1)变为了(n_heads, rows + 1)
            crow = crow[None].expand(n_heads, crow.shape[0])
            col = col[None].expand(n_heads, col.shape[0])
            if return_dense:
                # *是解包，拆分操作  shape为(5, 8) 拆分之后就是5, 8
                # 这里也是将expand操作，将单个token级掩码复制到n_heads上
                # 这个是在第0维度进行拓展操作 拓展到n_heads
                # 形状为(n_heads, *mask_dense.shape)
                mask_dense = mask_dense[None].expand(n_heads,
                                                     *mask_dense.shape)
                if dense_mask_type == "bias":
                    mask_dense = binary_mask_to_bias(mask_dense)
            return (crow, col), block_mask_dense, mask_dense

    # 非同质头操作
    with torch.no_grad():
        num_blocks = triton.cdiv(max_seqlen, block_size)
        # 这里进行了拓展维度
        # 形状 [1, num_blocks, 1]
        q_pos = torch.arange(num_blocks)[None, :, None]
        # 形状 [1, 1, num_blocks]
        k_pos = torch.arange(num_blocks)[None, None]
        # 获取滑动步长 均匀分配偏移量
        # 这一部分操作和上面的操作不同
        head_sliding_step = get_head_sliding_step(n_heads, vert_stride)
        # 为每个头h计算特定垂直步长掩码 生成n_heads个掩码列表
        mask_vert_strided = [
            (torch.arange(num_blocks) + h * head_sliding_step + 1) %
            vert_stride == 0 for h in range(n_heads)
        ]
        # 堆叠成(n_heads, 1, num_blocks)   unsqueeze(1) 相当于在第二个维度上增加一个维度
        mask_vert_strided = torch.vstack(mask_vert_strided).unsqueeze(1)
        block_mask_dense = (((q_pos >= k_pos)
                             & ((q_pos - k_pos < local_blocks)
                                | mask_vert_strided)).to(device).to(dtype))
        num_blocks_q = triton.cdiv(q_len, block_size)
        block_mask_dense_output = block_mask_dense[:, -num_blocks_q:]
    if return_dense:
        mask_dense = torch.kron(
            block_mask_dense,
            block_mask_dense.new_ones((block_size, block_size)),
        )
        causal_mask = torch.tril(torch.ones(
            max_seqlen, max_seqlen)).type_as(mask_dense)[-q_len:]
        mask_dense = mask_dense[..., -q_len:, :max_seqlen] * causal_mask[None]
        if dense_mask_type == "bias":
            mask_dense = binary_mask_to_bias(mask_dense)

        return (
            dense_to_crow_col(block_mask_dense_output),
            block_mask_dense,
            mask_dense,
        )
    else:
        return (
            dense_to_crow_col(block_mask_dense_output),
            block_mask_dense,
            None,
        )